#pragma once

#include <initializer_list>
#include <memory>
#include <string_view>
#include <variant>
#include <vector>

#include "kiva/common.h"

namespace kiva {

class Node;

// This class holds a raw pointer to a Node. Its initial usage is to work around
// the initializer_list's copy problem -- you can only copy things out of an
// initializer_list, so you cannot move Node out of it. Now we put NodeRef in
// initializer_list and when needed, own them with unique_ptr.
class NodeRef {
 public:
  NodeRef() {}
  NodeRef(Node* node) : node_(node) {}

  // Allow implicitly converted from Node.
  NodeRef(const std::unique_ptr<Node>& node) : node_(node.get()) {}

  // Trivially copy and move.
  NodeRef(NodeRef&&) = default;
  NodeRef(const NodeRef&) = default;
  NodeRef& operator=(NodeRef&&) = default;
  NodeRef& operator=(const NodeRef&) = default;

  // pointer null-check.
  explicit operator bool() const { return !!node_; }

  // Get the raw pointer, explicitly.
  Node* get() const { return node_; }

  // Call Node's methods.
  Node* operator->() const { return node_; }

  // When do we need a ref?
  Node& operator*() { return *node_; }
  const Node& operator*() const { return *node_; }

 private:
  Node* node_ = nullptr;
};

// Types of nodes.
// xxx: use the same set of names as Python.
enum class NodeType {
  kNop,

  // Expression
  kBoolOp,  // unpacked.
  kNamedExpr,
  kBinOP,    // unpacked.
  kUnaryOp,  // unpacked.
  kLambda,
  kIfExp,
  kDict,
  kSet,
  kListComp,
  kSetComp,
  kDictComp,
  kGeneratorExp,
  kAwait,
  kYield,
  kYieldFrom,
  kCompare,
  kCall,
  FormattedValue,  // nop
  JointStr,        // nop
  kConstant,
  kAttribute,
  kSubscript,
  kStarred,
  kName,
  kList,
  kTuple,

  // BinOp, keep the order!
  kAdd,
  kSub,
  kMult,
  kMatMult,  // nop
  kDiv,
  kMod,
  kPow,  // nop
  kLShift,
  kRShift,
  kBitOr,
  kBitXor,
  kBitAnd,
  kFloorDiv,  // nop
  // end of BinOP.

  // UnaryOp.
  kInvert,
  kNot,
  kUAdd,
  kUSub,

  // Our extension
  kUStar,
  kUAmpersand,
  // end of UnaryOp.

  // BoolOp.
  kAnd,
  kOr,
  // end of BoolOP.

  // Compare
  kEq,
  kNotEq,
  kLt,
  kLtE,
  kGt,
  kGtE,
  kIs,
  kIsNot,
  kIn,
  kNotIn,
  // end of Compare.

  // Our extensions to expr.
  // Constant
  kBigInt,
  kLong,
  // end of Constant.

  // end of Expression.

  // Slice
  kSlice,
  kExtSlice,
  kIndex,
  // end of slice.

  // Statement
  kFunctionDef,
  kAsyncFunctionDef,
  kClassDef,
  kReturn,
  kDelete,
  kAssign,
  kAugAssign,
  kFor,
  kAsyncFor,
  kWhile,
  kIf,
  kWith,
  kAsyncWith,
  kRaise,
  kTry,
  kAssert,
  kImport,
  kImportFrom,
  kGlobal,
  kNonLocal,
  kExpr,
  kPass,
  kBreak,
  kContinue,

  // Our extensions.
  kBody,
  kElse,
  kElif,
  kFinally,
  kExcept,
  // end of stmt.

  // Other stuffs.
  kAlias,
  kFormalArguments,

  kExpressionFirst = kBoolOp,
  kExpressionLast = kOr,

  kBinaryExprFirst = kAdd,
  kBinaryExprLast = kFloorDiv,

  kUnaryExprFirst = kInvert,
  kUnaryExprLast = kUAmpersand,

  kBoolOpFirst = kAnd,
  kBoolOpLast = kOr,

  kConstantFirst = kBigInt,
  kConstantLast = kLong,

  kCompareFirst = kEq,
  kCompareLast = kNotIn,

  kSliceFirst = kSlice,
  kSliceLast = kIndex,

  kStatementFirst = kFunctionDef,
  kStatementLast = kElif,

  kStatementNoChildFirst = kPass,
  kStatementNoChildLast = kContinue,
};

class Node {
 public:
  // ObjectType is used to store C++ constants that can be directly converted to
  // python object or just python object, like bool, double. And once they were
  // stored in the Node, they won't be accessed until being converted to
  // python'ast.
  using ObjectType = pybind11::object;
  // Internal rep of children nodes.
  using ListType = std::vector<std::unique_ptr<Node>>;
  // LongType is used to store integers that are not user specified constants
  // and can be read out fast from the Node.
  using LongType = long;
  // StringType is used to stored identifiers that are unchecked now and thus
  // not converted to python object.
  using StringType = std::string;

  // These are predicates about the category of a NodeType.
  static bool IsConstant(NodeType type);
  static bool IsBinOp(NodeType type);
  static bool IsUnaryOp(NodeType type);
  static bool IsBoolOp(NodeType type);
  static bool IsCompare(NodeType type);
  static bool IsExpression(NodeType type);
  static bool IsStatement(NodeType type);
  static bool IsStatementNoChild(NodeType type);
  static bool HasNoChild(NodeType type);

  // Default construct a Nop node.
  Node() {}

  // These constructors *have special purposes*. If in doubt, use named
  // constructors. Construct a node that can have children.
  Node(NodeType type, std::initializer_list<NodeRef> children);

  // Construct a node holding python object constants.
  explicit Node(ObjectType val);

  // Construct a node to store integer.
  explicit Node(LongType val);

  // Construct a node to store identifier.
  explicit Node(std::string_view val);

  // This class cannot be moved or copied. It can only be held by unique_ptr.
  //   Node(Node&&) = default;
  //   Node& operator=(Node&&) = default;

  NodeType type() const { return type_; }
  void set_type(NodeType type);

  // Access the underlying constant.
  ObjectType object_value() const;

  // Access the integer.
  long long_value() const;

  // Access the string/name/identifier.
  const StringType& string_value() const;
  // StringType& string_value();

  // Take ownership of this node.
  void AddChild(NodeRef node);
  void AddChild(std::unique_ptr<Node> node);

  ListType& children() { return as_list(); }
  const ListType& children() const { return as_list(); }

  Node* child_at(std::size_t i) { return as_list()[i].get(); }
  const Node* child_at(std::size_t i) const { return as_list()[i].get(); }

  std::size_t num_children() const { return as_list().size(); }

  // Named constructors:
  using Holder = std::unique_ptr<Node>;

  static NodeRef NewWithChildren(NodeType type,
                                 std::initializer_list<NodeRef> children = {});
  static NodeRef NewWithChildren(NodeType type, NodeRef child);

  static NodeRef NewName(std::string_view val);
  static NodeRef NewLong(LongType val);
  static NodeRef NewNodeType(NodeType val);

  static NodeRef NewConstant(long val);
  static NodeRef NewConstant(unsigned long val);
  static NodeRef NewConstant(double val);
  static NodeRef NewConstant(bool val);
  static NodeRef NewConstant(std::string_view val);

  static NodeRef NewBinOp(NodeRef lhs, NodeType type, NodeRef rhs);
  static NodeRef NewAssign(NodeRef lhs, NodeRef rhs);
  static NodeRef NewAugAssign(NodeRef lhs, NodeType type, NodeRef rhs);
  static NodeRef NewUnaryOp(NodeType type, NodeRef operand);
  static NodeRef NewBoolOp(NodeRef lhs, NodeType type, NodeRef rhs);
  static NodeRef NewCompare(NodeRef lhs, NodeType type, NodeRef rhs);
  static NodeRef NewAttribute(NodeRef lhs, std::string_view name);
  static NodeRef NewCall(NodeRef func, NodeRef args);
  static NodeRef NewSubscript(NodeRef node, NodeRef index);
  static NodeRef NewAlias(NodeRef name, std::string_view alias);

 private:
  ListType& as_list();
  const ListType& as_list() const;

  NodeType type_ = NodeType::kNop;
  std::variant<ObjectType, LongType, StringType, ListType> value_;
};

class NodeBuilder {
 public:
  // Start building an node from empty.
  explicit NodeBuilder(NodeType type);
  // Adopt an existing node.
  explicit NodeBuilder(NodeRef node);

  // Add children. Take ownership of children.
  NodeBuilder& AddChildren(std::initializer_list<NodeRef> children);
  NodeBuilder& AddChild(NodeRef child);
  NodeRef Build();

  // Return the node under construction.
  NodeRef node() const { return node_.get(); }

 private:
  std::unique_ptr<Node> node_;
};

}  // namespace kiva
